#!/usr/bin/env python
# coding=utf-8
import argparse
from collections import defaultdict, namedtuple
import functools
import sys
from time import sleep
from bcc import BPF
from socket import inet_ntop, AF_INET
from struct import pack
from subprocess import call

BPF_TXT = '''
#include <uapi/linux/ptrace.h>
#include <net/sock.h>
struct data_t{
    u32 pid;
    u32 addr;
    u16 sport;
    u16 dport;
};
BPF_HASH(tcp_send_bytes, struct data_t);
BPF_HASH(tcp_recv_bytes, struct data_t);
BPF_HASH(udp_send_bytes, struct data_t);
BPF_HASH(udp_recv_bytes, struct data_t);

int kprobe__tcp_sendmsg(struct pt_regs *ctx, struct sock *sk,
    struct msghdr *msg, size_t size)
{
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    ARG_PID

    u16 dport = 0, family = sk->__sk_common.skc_family;

    if (family == AF_INET) {
        struct data_t ipv4_key = {.pid=pid};
        ipv4_key.addr = sk->__sk_common.skc_daddr;
        ipv4_key.sport = sk->__sk_common.skc_num;
        dport = sk->__sk_common.skc_dport;
        ipv4_key.dport = ntohs(dport);
        tcp_send_bytes.increment(ipv4_key, size);
    }
    return 0;
}

int kprobe__tcp_recvmsg(struct pt_regs *ctx, struct sock *sk,  struct msghdr *msg, size_t len,int nonblock,int flags, int *addr_len)
{
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    ARG_PID

    u16 dport = 0, family = sk->__sk_common.skc_family;
    u64 *val, zero =0;

    if (len <= 0)
        return 0;

    if (family == AF_INET) {
        struct data_t ipv4_key = {.pid=pid};
        ipv4_key.addr = sk->__sk_common.skc_rcv_saddr;
        ipv4_key.sport = sk->__sk_common.skc_num;
        dport = sk->__sk_common.skc_dport;
        ipv4_key.dport = ntohs(dport);
        tcp_recv_bytes.increment(ipv4_key, len);
    }
    return 0;
}


int kprobe__udp_sendmsg(struct pt_regs *ctx, struct sock *sk,
    struct msghdr *msg, size_t size)
{
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    ARG_PID

    u16 dport = 0, family = sk->__sk_common.skc_family;

    if (family == AF_INET) {
        struct data_t ipv4_key = {.pid=pid};
        ipv4_key.addr = sk->__sk_common.skc_daddr;
        ipv4_key.sport = sk->__sk_common.skc_num;
        dport = sk->__sk_common.skc_dport;
        ipv4_key.dport = ntohs(dport);
        udp_send_bytes.increment(ipv4_key, size);
    }
    return 0;
}

int kprobe__udp_recvmsg(struct pt_regs *ctx, struct sock *sk,  struct msghdr *msg, size_t len,int nonblock,int flags, int *addr_len)
{
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    ARG_PID

    u16 dport = 0, family = sk->__sk_common.skc_family;
    u64 *val, zero =0;

    if (len <= 0)
        return 0;

    if (family == AF_INET) {
        struct data_t ipv4_key = {.pid=pid};
        ipv4_key.addr = sk->__sk_common.skc_rcv_saddr;
        ipv4_key.sport = sk->__sk_common.skc_num;
        dport = sk->__sk_common.skc_dport;
        ipv4_key.dport = ntohs(dport);
        udp_recv_bytes.increment(ipv4_key, len);
    }
    return 0;
}


'''


SessionKey = namedtuple('Session',[ 'pid','comm','sport', 'addr', 'dport',"type"])
@functools.cache
def pid_to_comm(pid):
    try:
        comm = open("/proc/%s/comm" % pid, "r").read().rstrip()
        return comm
    except IOError:
        return str(pid)

maxbit=sys.maxsize
x86 = True
if maxbit>2**32:
    x86 = False

def range_check(num)->bool:
    val = int(num)
    maxPid = 2**32 if x86 else 32768 
    if val<1 or val > maxPid:
        return False
    return True
def get_ipv4_session_key(k,isTcp):
    return SessionKey(pid=k.pid,comm=pid_to_comm(k.pid),sport=k.sport, addr=inet_ntop(AF_INET, pack("I", k.addr)), dport=k.dport,type="tcp" if isTcp else "udp")
def clear():
    call("clear")

parser = argparse.ArgumentParser(
    description = "Summarize send and recv flow by host",
    formatter_class = argparse.RawDescriptionHelpFormatter,
)
parser.add_argument("-p", "--pid",default="", 
    help = "Trace this pid only")
parser.add_argument("-i","--interval", default=1, type=int,
    help = "output interval, in second (default 1)")
arg = parser.parse_args()
def parser():
    global BPF_TXT
    if arg.pid !="":
        if range_check(arg.pid):
            Fliter_PID = f"if( pid != {arg.pid} ){{ return 0; }}"
    else:
            Fliter_PID = ""
    BPF_TXT = BPF_TXT.replace("ARG_PID",Fliter_PID)
    if arg.interval!=1:
        BPF_TXT = BPF_TXT.replace("ARG_INTERVAL",str(arg.interval))

if __name__ == "__main__":
    parser()
    b = BPF(text=BPF_TXT)

    tcp_send_bytes = b["tcp_send_bytes"]
    tcp_recv_bytes = b["tcp_recv_bytes"]
    udp_send_bytes = b["udp_send_bytes"]
    udp_recv_bytes = b["udp_recv_bytes"]

    # header
    print("%-10s %-12s %-10s %-10s %-10s %-10s %-10s %-21s %-21s %3s" % ("PID","COMM","RX_KB", "TX_KB", "RXSUM_KB", "TXSUM_KB", "SUM_KB", "LADDR", "RADDR","TYPE"))

    # output
    sumrecv = 0
    sumsend = 0
    sum_kb = 0
    i = 0
    exiting = False
    while i != 10 and not exiting:
        try:
            sleep(arg.interval)
        except KeyboardInterrupt:
            exiting = True

        ipv4_throughput = defaultdict(lambda:[0,0])
        for k, v in tcp_send_bytes.items():
            # print(k.addr,k.dport,k.sport)
            key=get_ipv4_session_key(k,True)
            ipv4_throughput[key][0] = v.value
        tcp_send_bytes.clear()

        for k,v in tcp_recv_bytes.items():
            key = get_ipv4_session_key(k,True)
            ipv4_throughput[key][1] = v.value
        tcp_recv_bytes.clear()
        for k, v in udp_send_bytes.items():
            # print(k.addr,k.dport,k.sport)
            key=get_ipv4_session_key(k,False)
            ipv4_throughput[key][0] = v.value
        udp_send_bytes.clear()

        for k,v in udp_recv_bytes.items():
            key = get_ipv4_session_key(k,False)
            ipv4_throughput[key][1] = v.value
        udp_recv_bytes.clear()
        #lock.acquire()
        res =""
        if ipv4_throughput:
            for k, (send_bytes, recv_bytes) in sorted(ipv4_throughput.items(),
                key=lambda kv: sum(kv[1]),
                reverse=True):
                recv_bytes = int(recv_bytes / 1024)
                send_bytes = int(send_bytes / 1024)
                if recv_bytes==0 and send_bytes==0:
                    continue
                sumrecv += recv_bytes
                sumsend += send_bytes
                sum_kb = sumrecv + sumsend
                res +="%-10d %-12.12s %-10d %-10d %-10d %-10d %-10d %-21s %-21s %3s\n" % (k.pid,k.comm,recv_bytes, send_bytes, sumrecv, sumsend, sum_kb, 
                    "localhost:" + str(k.sport), 
                    k.addr + ":" + str(k.dport),k.type)
        print(f"\r{res}",flush=True)
        #lock.release()
        i += 1